package cn.edu.flink.tutorial.state;

import cn.edu.flink.tutorial.source.UDSourceFunction;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.*;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

public class OperatorState {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Tuple2<Long, Long>> streamSource = env.addSource(new UDSourceFunction()).map(x -> Tuple2.of(x, 1L)).returns(Types.TUPLE(Types.LONG, Types.LONG));
        DataStreamSource<Long> elements = env.fromElements(1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L);

        // 一个 map descriptor，它描述了用于存储规则名称与规则本身的 map 存储结构
        MapStateDescriptor<Long, Long> mapStateDescriptor = new MapStateDescriptor<>("BroadcastState", Long.class, Long.class);
        BroadcastStream<Long> broadcast = elements.broadcast(mapStateDescriptor);

        streamSource
                .connect(broadcast)
                .process(new BroadcastProcessFunction<Tuple2<Long, Long>, Long, Long>() {
                    @Override
                    public void processElement(Tuple2<Long, Long> value, ReadOnlyContext ctx, Collector<Long> out) throws Exception {
                        System.out.println(value);
                        ReadOnlyBroadcastState<Long, Long> broadcast = ctx.getBroadcastState(mapStateDescriptor);
                        broadcast.immutableEntries().forEach(System.out::print);
                    }

                    @Override
                    public void processBroadcastElement(Long value, Context ctx, Collector<Long> out) throws Exception {
                        System.out.println("broadcast:" + value);
                        BroadcastState<Long, Long> ctxBroadcastState = ctx.getBroadcastState(mapStateDescriptor);
                         ctxBroadcastState.put(value, value);
                    }
                });


        env.execute();
    }
}
